import torch

"""
@网络模型：LeNet
@输入数据：batch_size*1*28*28灰度图像
@输出数据：10维度线性层分类结果
@学习率：0.9
"""


class LeNet(torch.nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self._pooling = torch.nn.AvgPool2d(kernel_size=2, stride=2)
        self._active = torch.nn.Sigmoid()
        self._flatten = torch.nn.Flatten()
        self._conv1 = torch.nn.Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1), padding=2)
        self._conv2 = torch.nn.Conv2d(6, 16, kernel_size=(5, 5), padding=1)
        self._linear1 = torch.nn.Linear(400, 120)
        self._linear2 = torch.nn.Linear(120, 84)
        self._linear3 = torch.nn.Linear(84, 10)

    def forward(self, x):
        x = x.view(-1, 1, 28, 28)
        x = self._conv1(x)
        x = self._pooling(self._active(x))
        x = self._conv2(x)
        x = self._pooling(self._active(x))
        x = self._flatten(x)
        x = self._active(self._linear1(x))
        x = self._active(self._linear2(x))
        return self._linear3(x)
